/* wave1d.c                                  (\(\
 * Jannis Teunnissen and Florian Speelman    (='.')
 * Scientific Computing - spring 2010        o(_")")
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <mpi.h>
#include <limits.h>
#include "divideElems.h"
#include "getWavePars.h"

#define WRITE_DATFILE 1
#define WRITE_PLOTS 2

void initPars(int argc, char *argv[], wavePars *p, MPI_Comm *ringPtr);
double waveSolver(wavePars *p, MPI_Comm *commPtr);
void initWave(double *u, wavePars *p, int *nElementsArray, int *displs, int myrank, int ntasks, MPI_Comm *ringPtr);
void writeWave(double *u, double time, wavePars *p, int mode, int *nElementsArray, int *displs, int myrank, MPI_Comm *ringPtr);
void stepWave(double *u, wavePars *p, int *nElemsArray, int *displs, int myrank, int ntasks, MPI_Comm *ringPtr);
double calcWaveError(double *uComplete, double time, wavePars *p);

int main(int argc, char *argv[]) {
    int rc;
    
    rc = MPI_Init(&argc, &argv);
    if(rc != MPI_SUCCESS) {
        printf("MPI_Init error %d\n", rc);
        MPI_Abort(MPI_COMM_WORLD, rc);
    }
    
    MPI_Comm ringComm;
    wavePars p;
    
    initPars(argc, argv, &p, &ringComm);
    
    double computeTime;
    computeTime = waveSolver(&p, &ringComm);
    
    if(computeTime > 0.) printf("Total compute time %e seconds\n", computeTime);
            
    MPI_Finalize();
    return 0;
}

double waveSolver(wavePars *p, MPI_Comm *ringPtr) {
    int myrank, ntasks;
    
    MPI_Comm_size(*ringPtr, &ntasks);
    MPI_Comm_rank(*ringPtr, &myrank);
    
    int nElemsArray[ntasks];    //Array containing the number of elements of each task,
    int displs[ntasks];         //Displacement array
    
    {
        int myWeight = 1;
        if(p->loadBalance == 1) {
            myWeight = benchmark();
        }
        int weights[ntasks];
        MPI_Gather(&myWeight, 1, MPI_INT, weights, 1, MPI_INT, 0, *ringPtr);
        if(myrank == 0) {
            printWeights(weights, ntasks);
            divideElems(nElemsArray, displs, weights, p->totalPoints, ntasks);
        }
    }
    
    //Now share nElemsArray and displs to later on use allgatherv
    MPI_Bcast(nElemsArray, ntasks, MPI_INT, 0, *ringPtr);
    MPI_Bcast(displs, ntasks, MPI_INT, 0, *ringPtr);
    int nPoints = nElemsArray[myrank];
    double *u = malloc((nPoints + 2) * sizeof(double));     //Displacement at t = 0

    initWave(u, p, nElemsArray, displs, myrank, ntasks, ringPtr);
        
    double startTime = 0., endTime = 0.;
    
    MPI_Barrier(*ringPtr);
    
    if(myrank == 0) startTime = MPI_Wtime();
    stepWave(u, p, nElemsArray, displs, myrank, ntasks, ringPtr);
    if(myrank == 0) endTime = MPI_Wtime();

    return endTime - startTime;
}

void stepWave(double *u, wavePars *p, int *nElemsArray, int *displs, int myrank, int ntasks, MPI_Comm *ringPtr) {
    int i, nPoints = nElemsArray[myrank];
    int left, right;
    double *u_old = malloc((nPoints + 2) * sizeof(double));
    double *u_new = malloc((nPoints + 2) * sizeof(double));
    double *tempPtr;
    double tau = p->c * p->c * p->dt * p->dt / (p->dx * p->dx);
    double twoMinTwoTau = 2. - 2. * tau;
    int outputCounter = 0;
    int startPos = 1, endPos = nPoints + 1;
    MPI_Cart_shift(*ringPtr, 0, 1, &left, &right);
    MPI_Request rq[4];
    
    if(p->fixedEnds == 1) {
        if(myrank == 0) {
            startPos = 2;
            left = MPI_PROC_NULL;
            u_new[1] = u_old[1] = u[1];
        }
        if(myrank == ntasks - 1){
            endPos = nPoints;
            right = MPI_PROC_NULL;
            u_new[nPoints] = u_old[nPoints] = u[nPoints];
        }
    }
    
    //Exchange halo cells for the first time
    MPI_Isend(&u[1], 1, MPI_DOUBLE, left, 0, *ringPtr, &rq[2]);
    MPI_Isend(&u[nPoints], 1, MPI_DOUBLE, right, 1, *ringPtr, &rq[3]);
    MPI_Irecv(&u[nPoints+1], 1, MPI_DOUBLE, right, 0, *ringPtr, &rq[0]);
    MPI_Irecv(&u[0], 1, MPI_DOUBLE, left, 1, *ringPtr, &rq[1]);
    MPI_Waitall(4, rq, MPI_STATUS_IGNORE);
    
    for(i = 0; i < nPoints + 2; i++) {
        u_old[i] = u[i];
    }

    int nTimeSteps = p->maxTime / p->dt;
    int timePerOutput;
    if(p->nOutputs > 0) {
        timePerOutput = nTimeSteps / p->nOutputs;
    } else {
        timePerOutput = nTimeSteps;         //No output
        outputCounter = 1;
    }

    for(int t = 0; t < nTimeSteps; t++) {
        if(t >= outputCounter * timePerOutput) {
            writeWave(u, p->dt * t, p, WRITE_DATFILE, nElemsArray, displs, myrank, ringPtr);
            outputCounter++;
        }
        u_new[startPos] = 2. * u[startPos] - u_old[startPos] + tau * (u[startPos+1] - 2. * u[startPos] + u[startPos-1]);
        u_new[endPos-1] = 2. * u[endPos-1] - u_old[endPos-1] + tau * (u[endPos] - 2. * u[endPos-1] + u[endPos-2]);
        
        MPI_Isend(&u_new[1], 1, MPI_DOUBLE, left, 0, *ringPtr, &rq[2]);
        MPI_Isend(&u_new[nPoints], 1, MPI_DOUBLE, right, 1, *ringPtr, &rq[3]);
        MPI_Irecv(&u_new[nPoints+1], 1, MPI_DOUBLE, right, 0, *ringPtr, &rq[0]);
        MPI_Irecv(&u_new[0], 1, MPI_DOUBLE, left, 1, *ringPtr, &rq[1]);
        
        for(i = startPos + 1; i < endPos - 1; i++) {
            u_new[i] = twoMinTwoTau * u[i] - u_old[i] + tau * (u[i+1] + u[i-1]);
        }

        MPI_Waitall(4, rq, MPI_STATUS_IGNORE);
        
        tempPtr = u_old;
        u_old = u;
        u = u_new;
        u_new = tempPtr;
    }
    
    if(p->imgOutput == 1) {
        writeWave(NULL, 0, p, WRITE_PLOTS, nElemsArray, displs, myrank, ringPtr);
    }
    free(u_new);
    free(u_old);
    free(u);
}

void writeWave(double *u, double time, wavePars *p, int mode, int *nElemsArray, int *displs, int myrank, MPI_Comm *ringPtr) {
    static int counter = 0;
    static double *uComplete = NULL;

    if(counter == 0 && myrank == 0) {
        FILE *datDir = fopen("dat/.dircheck", "w");
        FILE *pngDir = fopen("png/.dircheck", "w");
        FILE *outputInfo = fopen("dat/outputInfo.txt", "w");
        if(pngDir == NULL || datDir == NULL) {
            if(myrank == 0) fprintf(stderr, "Make sure directories 'dat' and 'png' exist\n");
            MPI_Abort(*ringPtr, 0);
        }
        fclose(datDir); fclose(pngDir); fclose(outputInfo);
        uComplete = malloc(p->totalPoints * sizeof(double));        //Don't free uComplete between calls for performance
    }
    
    if(mode == WRITE_DATFILE) {
        MPI_Gatherv(&u[1], nElemsArray[myrank], MPI_DOUBLE, uComplete, nElemsArray, displs, MPI_DOUBLE, 0, *ringPtr);
        if(myrank == 0) {
            char filename[40];
            sprintf(filename, "dat/wave_%d.dat", counter);
            FILE *output = fopen(filename, "w");
            FILE *outputInfo = fopen("dat/outputInfo.txt", "a");
            fprintf(outputInfo, "%e\t%e\n", time, calcWaveError(uComplete, time, p));
            for(int i = 0; i < p->totalPoints; i++) {
                fprintf(output, "%e\t%e\n", p->dx * i, uComplete[i]);
            }
            fclose(output); fclose(outputInfo);
            counter++;
        }
    }
    
    if (mode == WRITE_PLOTS && myrank == 0) {
        FILE *gnuplot = popen("gnuplot -persist", "w");
        FILE *outputInfo = fopen("dat/outputInfo.txt", "r");
        fprintf(gnuplot,"set nokey\n");
        fprintf(gnuplot,"set terminal png medium\n");
        
        fprintf(gnuplot,"set xlabel 'time'\n");
        fprintf(gnuplot,"set ylabel 'abs. error'\n");
        fprintf(gnuplot, "set title 'Absolute error over time'\n");
        fprintf(gnuplot, "set output 'png/waveError.png'\n");
        fprintf(gnuplot, "plot 'dat/outputInfo.txt' using 1:2 w l\n");
        
        fprintf(gnuplot,"set yrange [%lf:%lf]\n", -p->amplitude, p->amplitude);
        fprintf(gnuplot,"set xlabel 'x'\n");
        fprintf(gnuplot,"set ylabel 'u(x)'\n");

        double error;
        for(int i = 0; i < counter; i++) {
            if(fscanf(outputInfo, "%le %le\n", &time, &error)) {
                fprintf(gnuplot, "set title 'String at t = %.2lf s, max error %.2e'\n", time, error);
                fprintf(gnuplot, "set output 'png/wave_%06d.png'\n", i);
                fprintf(gnuplot, "plot 'dat/wave_%d.dat' using 1:2 w l\n", i);
            }
        }
        fclose(gnuplot); fclose(outputInfo);
    }
}       
    

void initWave(double *u, wavePars *p, int *nElemsArray, int *displs, int myrank, int ntasks, MPI_Comm *ringPtr) {
    int i;
    double *uComplete = NULL;
    
    if(myrank == 0) {
        uComplete = malloc((p->totalPoints) * sizeof(double));
        
        if(p->plucked == 1) {
            int pluckLocation = round(p->location / p->dx);
            double pluckFactorLeft = p->dx * p->amplitude / p->location;
            double pluckFactorRight = p->amplitude / (1. - p->location);
            
            for(i = 0; i < pluckLocation && i < p->totalPoints; i++) {
                uComplete[i] = pluckFactorLeft * i;
            }
            for(i = pluckLocation; i < p->totalPoints; i++) {
                uComplete[i] = pluckFactorRight * (1. - p->dx * i);
            }
        } else {
            double sinFactor = 2. * acos(-1.) / p->sinPeriod;
            for(i = 0; i < p->totalPoints; i++) {
                uComplete[i] = p->amplitude * sin(p->dx * i * sinFactor);
            }
        }
    }
    
    MPI_Scatterv(uComplete, nElemsArray, displs, MPI_DOUBLE, &u[1], nElemsArray[myrank], MPI_DOUBLE, 0, *ringPtr);
    if(myrank == 0) {
        free(uComplete);
    }
}

double calcWaveError(double *uComplete, double time, wavePars *p) {
    double sol, maxdiff, min, max, diff, x, x1, x2, temp;
    double sinFactor = 2. * acos(-1.) / p->sinPeriod;
    maxdiff = 0.;
    
    for(int i = 0; i < p->totalPoints; i++) {
        min = i * p->dx - p->c * time;
        max = i * p->dx + p->c * time;
        
        if(p->plucked == 0) {
            sol = 0.5 * p->amplitude * (sin(max * sinFactor) + sin(min * sinFactor));
        } else {
            sol = 0.;
            x1 = fmod(fmod(min, 2.) + 2., 2.);
            x2 = fmod(max, 2.);
            
            x = fmin(x1, 2.-x1);
            if(x < p->location) {
                temp = 0.5 * p->amplitude / p->location * x;
            } else {
                temp = 0.5 * p->amplitude / (1. - p->location) * (1. - x);
            }
            if(x1 > 1.) temp *= -1.;
            sol += temp;
            
            x = fmin(x2, 2.-x2);
            if(x < p->location) {
                temp = 0.5 * p->amplitude / p->location * x;
            } else {
                temp = 0.5 * p->amplitude / (1. - p->location) * (1. - x);
            }
            if(x2 > 1.) temp *= -1.;
            sol += temp;
        }
        diff = fabs(uComplete[i] - sol);
        maxdiff = diff > maxdiff ? diff : maxdiff;
    }
    return maxdiff;
}
            

void initPars(int argc, char *argv[], wavePars *p, MPI_Comm *ringPtr) {
    int ntasks, myrank, one = 1;
    
    MPI_Comm_size(MPI_COMM_WORLD, &ntasks);
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
    
    MPI_Cart_create(MPI_COMM_WORLD, one, &ntasks, &one, one, ringPtr);
    
    if(myrank == 0) getWavePars(argc, argv, p);

    MPI_Datatype wavepMPI;
    MPI_Datatype ptype[2] = {MPI_INT, MPI_DOUBLE};
    MPI_Aint pdisp[2] = {(void*)&p->plucked - (void*)p, (void*)&p->dx - (void*)p};
    int pblocklen[2] = {6, 7};
    MPI_Type_create_struct(2, pblocklen, pdisp, ptype, &wavepMPI);
    
    MPI_Type_commit(&wavepMPI);
    MPI_Bcast(p, 1, wavepMPI, 0, *ringPtr);
    MPI_Type_free(&wavepMPI);
}


